import os
import cv2
import pynvml
import warnings
#! 设置显卡
pynvml.nvmlInit() #! 初始化
deviceCount = pynvml.nvmlDeviceGetCount() #! 统计显卡数量
GPUs = []
for busid in range(deviceCount):
    handle = pynvml.nvmlDeviceGetHandleByIndex(busid)
    device = pynvml.nvmlDeviceGetMemoryInfo(handle)
    memory = device.total // 1048576 / 1024 #! 单位: G
    if memory > 8: GPUs.append(str(busid))
pynvml.nvmlShutdown() #! 关闭
GPUs = []
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(GPUs) #! 设置显卡设备
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" #! 设置警告ERROR级别
#! 导入官方包
import numpy as np
import tensorflow as tf
import tensorflow.keras as keras
import tensorflow.keras.backend as K
#! 导入自定义模块
from model import SSD300, SSD512
from data import PascalVOCDatasetV1, PascalVOCDatasetV2
#! 忽略警告类报错
warnings.filterwarnings('ignore')

#! 数据集设置
data_folder = 'dataset/VOCdevkit'
voc_labels = ('aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable',
              'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor')
label_map = {k: v + 1 for v, k in enumerate(voc_labels)}
label_map['background'] = 0
n_classes = len(label_map)
rev_label_map = {v:k for k,v in label_map.items()}
batch_size = 16
PascalVOCDataset = PascalVOCDatasetV2
output_shape = (300,300)
# train_loader = PascalVOCDataset(batch_size,data_folder,'TRAIN',output_shape=output_shape,shuffle=True)
# valid_loader = PascalVOCDataset(batch_size,data_folder,'TEST',output_shape=output_shape)

#! 模型设置
model_path = None
SSD = SSD300
SSDModel = SSD(n_classes=n_classes, model_path=model_path) if model_path is not None else SSD(n_classes=n_classes)

def main():
    model = SSDModel.model

    data_loader = PascalVOCDataset(batch_size,data_folder,'TEST',output_shape=output_shape, transform=False, shuffle=True)
    
    for i, (images, batch_image_boxes, batch_image_labels, _) in enumerate(data_loader):

        #! 原始图片
        batch_image_paths = data_loader.images[i * batch_size: (i+1) * batch_size]

        #! 预测结果
        predicted_locs, predicted_scores = model(images)
        batch_image_boxes, batch_image_labels, batch_image_scores =\
            SSDModel.detect_objects(predicted_locs, predicted_scores, min_score=0.05, max_overlap=0.6, top_k=100)

        for j,(image_path, image_boxes, image_labels, image_scores) in \
            enumerate(zip(batch_image_paths,batch_image_boxes, batch_image_labels, batch_image_scores)):

            image = cv2.imread(image_path)
            h, w, c = image.shape

            image_boxes = (image_boxes * np.array([w,h,w,h])).astype(np.int64)
            for c in range(1,n_classes):
                indices, = np.where(image_labels==c)
                if not indices.any(): continue

                class_image_boxes  = image_boxes [indices]
                class_image_scores = image_scores[indices]

                for box, score in zip(class_image_boxes, class_image_scores):
                    if score > 0.45:
                        x1, y1, x2, y2 = box
                        image = cv2.rectangle(image,(x1,y1),(x2,y2),(255,0,0))
                        print(rev_label_map[int(c)], score)
            
            print("-=-"*30)
            
            cv2.imshow('image',image)
            cv2.waitKey()
            cv2.destroyAllWindows()

            # break

        break


if __name__ == '__main__':
    main()
